/*
Copyright 2024 MySQL Parser Project

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package sqlparser

import (
	"strings"
)

// extractTablesFromTableExpr extracts table names from a TableExpr
func extractTablesFromTableExpr(expr TableExpr) []string {
	var tables []string
	tableMap := make(map[string]bool)

	Walk(func(node SQLNode) (bool, error) {
		switch n := node.(type) {
		case *TableName:
			if !n.IsEmpty() {
				tableName := n.Name.String()
				if !tableMap[tableName] {
					tables = append(tables, tableName)
					tableMap[tableName] = true
				}
			}
		case TableName:
			if !n.IsEmpty() {
				tableName := n.Name.String()
				if !tableMap[tableName] {
					tables = append(tables, tableName)
					tableMap[tableName] = true
				}
			}
		}
		return true, nil
	}, expr)

	return tables
}

// ExtractAllTables extracts all table names from a SQL statement
// This function is inspired by Vitess's ExtractAllTables function
func ExtractAllTables(stmt Statement) []string {
	var tables []string
	tableMap := make(map[string]bool)

	// Use AST walking to find all table references
	Walk(func(node SQLNode) (bool, error) {
		switch n := node.(type) {
		case *TableName:
			if !n.IsEmpty() {
				tableName := n.Name.String()
				if !tableMap[tableName] {
					tables = append(tables, tableName)
					tableMap[tableName] = true
				}
			}
		case TableName:
			if !n.IsEmpty() {
				tableName := n.Name.String()
				if !tableMap[tableName] {
					tables = append(tables, tableName)
					tableMap[tableName] = true
				}
			}
		case *AliasedTableExpr:
			if tableName, ok := n.Expr.(*TableName); ok && !tableName.IsEmpty() {
				name := tableName.Name.String()
				if !tableMap[name] {
					tables = append(tables, name)
					tableMap[name] = true
				}
			} else if tableName, ok := n.Expr.(TableName); ok && !tableName.IsEmpty() {
				name := tableName.Name.String()
				if !tableMap[name] {
					tables = append(tables, name)
					tableMap[name] = true
				}
			} else if subquery, ok := n.Expr.(*Subquery); ok {
				// Recursively extract tables from subquery
				subTables := ExtractAllTables(subquery.Select)
				for _, subTable := range subTables {
					if !tableMap[subTable] {
						tables = append(tables, subTable)
						tableMap[subTable] = true
					}
				}
			}
		case *JoinTableExpr:
			// Extract tables from both sides of the join
			leftTables := extractTablesFromTableExpr(n.LeftExpr)
			rightTables := extractTablesFromTableExpr(n.RightExpr)
			for _, table := range append(leftTables, rightTables...) {
				if !tableMap[table] {
					tables = append(tables, table)
					tableMap[table] = true
				}
			}
		case *Subquery:
			// Extract tables from subquery
			subTables := ExtractAllTables(n.Select)
			for _, subTable := range subTables {
				if !tableMap[subTable] {
					tables = append(tables, subTable)
					tableMap[subTable] = true
				}
			}
		}
		return true, nil
	}, stmt)

	return tables
}

// ContainsAggregation checks if a SQL statement contains aggregation functions
// This function is inspired by Vitess's ContainsAggregation function
func ContainsAggregation(stmt Statement) bool {
	hasAggregation := false

	// Use AST walking to find aggregation functions
	Walk(func(node SQLNode) (bool, error) {
		switch n := node.(type) {
		case *FuncExpr:
			// Check if it's an aggregation function
			funcName := strings.ToLower(n.Name.String())
			if isAggregationFunction(funcName) {
				hasAggregation = true
				return false, nil // Stop walking
			}
		case *Select:
			// Check for GROUP BY clause
			if len(n.GroupBy) > 0 {
				hasAggregation = true
				return false, nil // Stop walking
			}
			// Check for HAVING clause
			if n.Having != nil {
				hasAggregation = true
				return false, nil // Stop walking
			}
		}
		return true, nil
	}, stmt)

	return hasAggregation
}

// isAggregationFunction checks if a function name is an aggregation function
func isAggregationFunction(name string) bool {
	lowerName := strings.ToLower(name)
	aggregationFunctions := map[string]bool{
		"count":     true,
		"sum":       true,
		"avg":       true,
		"min":       true,
		"max":       true,
		"group_concat": true,
		"bit_and":   true,
		"bit_or":    true,
		"bit_xor":   true,
		"std":       true,
		"stddev":    true,
		"stddev_pop": true,
		"stddev_samp": true,
		"var_pop":   true,
		"var_samp":  true,
		"variance":  true,
		"json_arrayagg": true,
		"json_objectagg": true,
	}
	return aggregationFunctions[lowerName]
}

// IsAggregateFunction checks if a function name is an aggregate function
func IsAggregateFunction(funcName string) bool {
	funcName = strings.ToUpper(funcName)
	aggregateFunctions := map[string]bool{
		"COUNT":     true,
		"SUM":       true,
		"AVG":       true,
		"MIN":       true,
		"MAX":       true,
		"GROUP_CONCAT": true,
		"BIT_AND":   true,
		"BIT_OR":    true,
		"BIT_XOR":   true,
		"STD":       true,
		"STDDEV":    true,
		"STDDEV_POP": true,
		"STDDEV_SAMP": true,
		"VAR_POP":   true,
		"VAR_SAMP":  true,
		"VARIANCE":  true,
		"JSON_ARRAYAGG": true,
		"JSON_OBJECTAGG": true,
	}
	return aggregateFunctions[funcName]
}



// TableAnalysis represents the result of table analysis
type TableAnalysis struct {
	Tables      []string
	HasJoins    bool
	HasSubquery bool
	TableCount  int
}

// AnalyzeTables analyzes tables in a SQL statement
func AnalyzeTables(stmt Statement) *TableAnalysis {
	tables := ExtractAllTables(stmt)
	analysis := &TableAnalysis{
		Tables:     tables,
		TableCount: len(tables),
	}

	// Check for joins and subqueries
	Walk(func(node SQLNode) (bool, error) {
		switch n := node.(type) {
		case *Select:
			if len(n.From) > 1 {
				analysis.HasJoins = true
			}
			// Check for subqueries in FROM clause
			for _, tableExpr := range n.From {
				if aliased, ok := tableExpr.(*AliasedTableExpr); ok {
					if _, isSubquery := aliased.Expr.(*Subquery); isSubquery {
						analysis.HasSubquery = true
					}
				}
			}
			// Check for subqueries in WHERE clause
			if n.Where != nil {
				Walk(func(subNode SQLNode) (bool, error) {
					if _, isSubquery := subNode.(*Subquery); isSubquery {
						analysis.HasSubquery = true
					}
					return true, nil
				}, n.Where)
			}
		case *JoinTableExpr:
			analysis.HasJoins = true
		case *Subquery:
			analysis.HasSubquery = true
		}
		return true, nil
	}, stmt)

	return analysis
}

// AggregationAnalysis represents the result of aggregation analysis
type AggregationAnalysis struct {
	HasAggregation bool
	Functions      []string
	HasGroupBy     bool
	HasHaving      bool
}

// AnalyzeAggregation analyzes aggregation in a SQL statement
func AnalyzeAggregation(stmt Statement) *AggregationAnalysis {
	analysis := &AggregationAnalysis{
		HasAggregation: false,
		Functions:      []string{},
	}

	funcMap := make(map[string]bool)

	Walk(func(node SQLNode) (bool, error) {
		switch n := node.(type) {
		case *FuncExpr:
			funcName := strings.ToUpper(n.Name.String())
			if isAggregationFunction(strings.ToLower(funcName)) {
				if !funcMap[funcName] {
					analysis.Functions = append(analysis.Functions, funcName)
					funcMap[funcName] = true
				}
				analysis.HasAggregation = true
			}
		case *Select:
			if len(n.GroupBy) > 0 {
				analysis.HasGroupBy = true
				analysis.HasAggregation = true
			}
			if n.Having != nil {
				analysis.HasHaving = true
				analysis.HasAggregation = true
			}
		}
		return true, nil
	}, stmt)

	return analysis
}

// QueryComplexity represents the complexity analysis of a query
type QueryComplexity struct {
	TableAnalysis       *TableAnalysis
	AggregationAnalysis *AggregationAnalysis
	ComplexityScore     int
	ComplexityLevel     string
}

// AnalyzeQueryComplexity performs comprehensive query analysis
func AnalyzeQueryComplexity(stmt Statement) *QueryComplexity {
	tableAnalysis := AnalyzeTables(stmt)
	aggrAnalysis := AnalyzeAggregation(stmt)

	// Calculate complexity score
	score := 0
	score += tableAnalysis.TableCount * 2
	if tableAnalysis.HasJoins {
		score += 5
	}
	if tableAnalysis.HasSubquery {
		score += 3
	}
	if aggrAnalysis.HasAggregation {
		score += 4
	}
	if aggrAnalysis.HasGroupBy {
		score += 2
	}
	if aggrAnalysis.HasHaving {
		score += 2
	}

	// Determine complexity level
	level := "Simple"
	if score >= 15 {
		level = "Very Complex"
	} else if score >= 10 {
		level = "Complex"
	} else if score >= 5 {
		level = "Moderate"
	}

	return &QueryComplexity{
		TableAnalysis:       tableAnalysis,
		AggregationAnalysis: aggrAnalysis,
		ComplexityScore:     score,
		ComplexityLevel:     level,
	}
}